import cv2
import numpy as np
import yaml
from utils.reader import from_TIFF, from_IMG

config_path = "/home/a804_cbf/Code/luner_crater/config/2080x4/HG-HV.yaml"
with open(config_path, "r") as f:
    config = yaml.load(f, Loader=yaml.FullLoader)
    nac_config = config["nac"]
    dem_config = config["dem"]
    subfig_config = config["subfig"]
dem_prj = from_TIFF(**dem_config, **subfig_config)
nac_prj = from_IMG(nac_config, **subfig_config)
nac = nac_prj.image.copy()
dem = dem_prj.image.copy()
nac = (nac - nac.min()) / (nac.max() - nac.min()) * 255
dem[dem == dem.min()] = dem.max()
dem = (dem - dem.min()) / (dem.max() - dem.min()) * 255
nac = cv2.cvtColor(nac, cv2.COLOR_GRAY2BGR)
dem = cv2.cvtColor(dem, cv2.COLOR_GRAY2BGR)
with open(
    "/disk527/sdb1/a804_cbf/datasets/chang_e/crater_catalog.txt", "r"
) as f:
    for line in f:
        lat, lon, r = map(float, line.split(","))
        x1, y1 = nac_prj.forward(lat, lon)
        x2, y2 = dem_prj.forward(lat, lon)
        cv2.circle(nac, (round(y1), round(x1)), round(r), (0, 0, 255), 1)
        cv2.ellipse(
            dem,
            (round(y2), round(x2)),
            (round(r), round(r * np.cos(lat / 180 * np.pi))),
            0,
            0,
            360,
            (0, 0, 255),
            1,
        )
cv2.imwrite("nac-craters.png", nac)
cv2.imwrite("dem-craters.png", dem)
